# data:下载放入 主目录下的datasets内。链接：https://pan.baidu.com/s/1X3tf2An0fTpGpZwjlYeglQ 提取码：aiAI
# From
# pytorch教程：https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
# pytorch代码：https://github.com/pytorch/tutorials/blob/master/beginner_source/blitz/cifar10_tutorial.py

import torch
import torchvision
import torch.nn as nn
from network import LeNet
import torch.optim as optim
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np
import argparse


def parse_option():
    parser = argparse.ArgumentParser('LeNet')
    parser.add_argument('--epochs', type=int, help='the number of epoch', default=4)
    parser.add_argument('--batch_size', type=int, help='batch_size', default=32)
    parser.add_argument("--data_path", type=str, help="the path of model saved",
                        default='../../datasets/CIFAR10')
    parser.add_argument("--save_path", type=str, help="the path of model saved",
                        default='../../models/LeNet.pth')
    args = parser.parse_args()  # 也可直接使用 args, _ = parser.parse_known_args()
    return args


def main(args):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("using {} device.".format(device))
    
    transform = transforms.Compose(
        [transforms.ToTensor(),
         transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

    # CIFAR10 50000张训练图片，10000张验证图片
    # 注意，首次使用：注意 download=True
    train_set = torchvision.datasets.CIFAR10(root=args.data_path, train=True,
                                             download=False, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size,
                                               shuffle=True, num_workers=0)
    # 注意，首次使用：注意 download=True
    test_set = torchvision.datasets.CIFAR10(root=args.data_path, train=False,
                                            download=False, transform=transform)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.batch_size,
                                              shuffle=False, num_workers=0)
    print(f"using {len(train_set)}  for training, {len(test_set)}  for validation.")


    # classes = ('plane', 'car', 'bird', 'cat',
    #            'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    # def imshow(img):
    #     img = img / 2 + 0.5  # unnormalize
    #     npimg = img.numpy()
    #     plt.imshow(np.transpose(npimg, (1, 2, 0)))  # 转成【高,宽,通道】的格式
    #     plt.show()
    # # get some random training images
    # dataiter = iter(train_loader)
    # images, labels = next(dataiter)
    # # show images
    # imshow(torchvision.utils.make_grid(images))
    # # print labels
    # print(' '.join(f'{classes[labels[j]]:5s}' for j in range(args.batch_size)))


    net = LeNet()
    net.to(device)
    criterion = nn.CrossEntropyLoss()   # loss function
    optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

    for epoch in range(args.epochs):  # loop over the dataset multiple times
        net.train()
        running_loss = 0.0
        for i, data in enumerate(train_loader, start=0):
            # get the inputs; data is a list of [inputs, labels]
            inputs, labels = data

            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs.to(device))
            loss = criterion(outputs, labels.to(device))
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            if i % 2000 == 1999:  # print every 2000 mini-batches
                print(f'train_data:[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
                running_loss = 0.0

        correct = 0
        total = 0
        # since we're not training, we don't need to calculate the gradients for our outputs
        net.eval()
        with torch.no_grad():
            for data in test_loader:
                outputs, test_labels = data
                # calculate outputs by running images through the network
                outputs = net(outputs.to(device))
                # the class with the highest energy is what we choose as prediction
                _, predict_y = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += torch.eq(predict_y, test_labels.to(device)).sum().item()
        print(f'Accuracy of the network on the 10000 test images: {100 * correct // total} %')

    print('Finished Training')
    torch.save(net.state_dict(), args.save_path)


if __name__ == '__main__':
    args_ = parse_option()
    main(args_)
